import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
import joblib

# 1. 读取训练数据（确保列名正确）
train_data = pd.read_csv("任务四词语.csv", encoding='utf-8')  # 假设使用|分隔符

# 提取文本（X）和类别（y）
texts = train_data["词语"]    # 输入特征
labels = train_data["类别"]   # 输出标签

# 2. TF-IDF 向量化
vectorizer = TfidfVectorizer()
X_train = vectorizer.fit_transform(texts)

# 3. 训练模型（逻辑回归）
model = LogisticRegression(
    class_weight='balanced',
    multi_class='multinomial',
    solver='lbfgs',
    max_iter=1000
)
model.fit(X_train, labels)

# 4. 保存模型和转换器
joblib.dump(model, "../hexinwenjianbao/task4_classifier.pkl")      # 保存模型
joblib.dump(vectorizer, "../hexinwenjianbao/task4_vectorizer.pkl")# 保存文本转数字的工具

print("模型训练完成！")

